#!/usr/bin/python
# _*_ coding: utf-8 _*_

"""
出填空题的样本

@author: Qing Liu, liuqing1@aircas.ac.cn
@time: 2024/5/17 9:01
"""
import json
import requests
from transformers import AutoTokenizer, AutoModel
import os
import re
from flask import request, stream_with_context, jsonify, Blueprint

import random
import math

chatglm_bp = Blueprint('chatglm', __name__)

model_path = os.path.join('/', 'usr', 'src', 'ZhipuAI', 'chatglm3-6b')
print(f'model_path is {model_path}', flush=True)

with open('entity.json', encoding='utf8') as fin:
    entities = json.load(fin)['RECORDS']
    print(len(entities))

fout = open('chat_qa_samples.json', 'w', encoding='utf8')

tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map='auto')
model = model.eval()


def ask_llm(prompt):
    answer, history = model.chat(tokenizer, prompt, history=[])
    return json.loads(answer)


for entity in entities:
    content = entity['JS']
    print("文本上下文：", content)
    try:
        # 出题：
        prompt = """
        你是一个专业的填空题出题老师，基于给定的文本上下文，出3-5个填空题，其中抠选出信息是该文本片段中包含的关键信息，包括时间、地点、组织机构、关键词、人物、武器装备、基地等重要信息。
        返回标准的json格式，注意是多个填空题样本，因此返回数组类型，格式如下：[{"问题": question, "答案": answer}]
        
        例如：
        文本上下文：日向级护卫舰（ひゅうが型護衛艦）是日本海上自卫队旗下的一个直升机护卫舰（Helicopter Destroyer, DDH）舰级。本级舰共两艘，首舰日向号于2006开工，2009年3月18日就役。?二号舰伊势号，这一舰名源自日本古代的伊势国，即现在的三重县。于2011年3月16日就役。
        填空题：
        [
            {
                "问题": "日向级护卫舰是（）旗下的一个直升机护卫舰舰级。",
                "答案": ["日本海上自卫队"]
            }，
            {
                "问题": "日向级护卫舰本级舰共两艘，首舰（）于2006开工，2009年3月18日就役，二号（），于2011年3月16日就役。",
                "答案": ["日向号", "舰伊势号"]
            }
        ]
        
        文本上下文：{context}
        填空题：
        """.replace("{context}", content)
        qa_sample = ask_llm(prompt)
        qa_sample = json.loads(qa_sample)

        # 答题
        answer_prompt_template = """
        你是一个专业的填空题答题学生，基于给定的文本上下文，回答填空题，请直接给出答案，返回答案必须是数组类型，因为可能会存在多个答案。

        例如：
        文本上下文：卡-26/126是为农业和其他空中作业研制的多功能直升机，1965年首飞。引人注目的是它的驾驶舱后面安装了可互换模块，可针对客运/货运、空中救护、消防及其他用途而进行调换。1970年，卡-26/126直升机在前苏联开始了它的服役生涯。该机型共计交付了超过850架。
        问题：卡-26/126是为（）和（）研制的多功能直升机。
        答案：["农业", "空中作业"]
        问题：1970年，卡-26/126直升机在（）开始了它的服役生涯。
        答案：["前苏联"]

        文本上下文：{context}
        问题：{question}
        答案：
        """
        for qa in qa_sample:
            answer_prompt = answer_prompt_template.replace('{context}', content) \
                .replace('{question}', qa['问题'])
            answer = ask_llm(answer_prompt)
            print('问题：' + qa['问题'])
            print('答案：' + json.dumps(qa['答案'], ensure_ascii=False))
            print('答题：' + answer)

            qa['答题'] = json.loads(answer)

        qa_sample = {
            "文本上下文": content,
            "填空题": qa_sample
        }
    except Exception as e:
        print(e, flush=True)
        continue

    fout.write(json.dumps(qa_sample, ensure_ascii=False, indent=4) + '\n')
